{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8a6826e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import random\n",
    "import shutil\n",
    "import time\n",
    "\n",
    "import torch\n",
    "import torch.backends.cudnn as cudnn\n",
    "import torch.nn as nn\n",
    "import torch.nn.parallel\n",
    "import torch.optim\n",
    "import torch.utils.data\n",
    "import torchvision.datasets as datasets\n",
    "import torchvision.models as models\n",
    "import torchvision.transforms as transforms\n",
    "import torch\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddfa2a76-7112-4860-91ed-ca7525025db8",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27d6825b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "is_CELU = False\n",
    "is_ELU = False\n",
    "\n",
    "\n",
    "#SE模块\n",
    "class SEModule(nn.Module):\n",
    "    def __init__(self, channels, reduction=16):\n",
    "        super(SEModule, self).__init__()\n",
    "        self.avg_pool = nn.AdaptiveAvgPool2d(1)\n",
    "        self.fc1 = nn.Conv2d(channels, channels // reduction, kernel_size=1, padding=0)\n",
    "        if is_CELU :\n",
    "            self.af = nn.CELU(inplace=True)\n",
    "        elif is_ELU :\n",
    "            self.af = nn.ELU(inplace=True)\n",
    "        else :\n",
    "            self.af = nn.ReLU(inplace=True)\n",
    "        self.fc2 = nn.Conv2d(channels // reduction, channels, kernel_size=1, padding=0)\n",
    "        self.sigmoid = nn.Sigmoid()\n",
    "\n",
    "    def forward(self, input):\n",
    "        x = self.avg_pool(input)\n",
    "        x = self.fc1(x)\n",
    "        #x = self.relu(x)\n",
    "        x = self.af(x)\n",
    "        x = self.fc2(x)\n",
    "        x = self.sigmoid(x)\n",
    "        return input * x\n",
    "\n",
    "\n",
    "# 残差层\n",
    "class Bottleneck(nn.Module):\n",
    "    # 输出通道数为输入通道数的4倍\n",
    "    expansion = 4\n",
    "\n",
    "    def __init__(self, inplanes, planes, stride=1, downsample=None, scales=4, groups=1, is_first_block=0, se=True):\n",
    "        super(Bottleneck, self).__init__()\n",
    "\n",
    "        self.downsample = downsample\n",
    "        self.scales = scales\n",
    "        self.groups = groups\n",
    "        self.stride = stride\n",
    "        self.is_first_block = is_first_block\n",
    "\n",
    "        outplanes = groups * planes\n",
    "        # 第一个卷积层：卷积核尺寸： 1×1 ，填充值为 0， 步长为 1\n",
    "        self.conv1 = nn.Conv2d(in_channels=inplanes, out_channels=outplanes, kernel_size=1, stride=1, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(outplanes)\n",
    "\n",
    "        # 第二个卷积结构：卷积核尺寸： 3×3 ，填充值为 1， 步长为 1\n",
    "        self.conv2 = nn.ModuleList([nn.Conv2d(outplanes // scales, outplanes // scales,\n",
    "                                              kernel_size=3, stride=stride, padding=1, groups=groups, bias=False) for _ in\n",
    "                                    range(scales - 1)])\n",
    "        self.bn2 = nn.ModuleList([nn.BatchNorm2d(outplanes // scales) for _ in range(scales - 1)])\n",
    "\n",
    "        # 第三个卷积层：卷积核尺寸： 1×1 ，填充值为 0， 步长为 1\n",
    "        self.conv3 = nn.Conv2d(outplanes, planes * self.expansion, kernel_size=1, stride=1, bias=False)\n",
    "        self.bn3 = nn.BatchNorm2d(planes * self.expansion)\n",
    "\n",
    "        #self.relu = nn.ReLU(inplace=True)\n",
    "        if is_CELU :\n",
    "            self.af = nn.CELU(inplace=True)\n",
    "        elif is_ELU :\n",
    "            self.af = nn.ELU(inplace=True)\n",
    "        else :\n",
    "            self.af = nn.ReLU(inplace=True)\n",
    "        \n",
    "        # 处理第一个块\n",
    "        if is_first_block == 1:\n",
    "            self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)\n",
    "\n",
    "        # SE模块\n",
    "        self.se = SEModule(planes * self.expansion) if se else None\n",
    "\n",
    "    def forward(self, x):\n",
    "        identity = x # 将原始输入暂存为shortcut的输出\n",
    "        # 对下采样进行处理\n",
    "        if self.downsample is not None:\n",
    "            identity = self.downsample(identity)\n",
    "\n",
    "        # 1*1卷积层\n",
    "        out = self.conv1(x)\n",
    "        out = self.bn1(out)\n",
    "        #out = self.relu(out)\n",
    "        out = self.af(out)\n",
    "\n",
    "        x_scales = torch.chunk(out, self.scales, 1)\n",
    "        for i in range(self.scales-1):\n",
    "            if i == 0 or self.is_first_block == 1:\n",
    "                y_scale = x_scales[i]\n",
    "            else:\n",
    "                y_scale = y_scale + x_scales[i]\n",
    "            y_scale = self.conv2[i](y_scale)\n",
    "            #y_scale = self.relu(self.bn2[i](y_scale))\n",
    "            y_scale = self.af(self.bn2[i](y_scale))\n",
    "            if i == 0:\n",
    "                out = y_scale\n",
    "            else:\n",
    "                out = torch.cat((out, y_scale), 1)\n",
    "        if self.scales != 1 and self.is_first_block == 0:\n",
    "            out = torch.cat((out, x_scales[self.scales-1]), 1)\n",
    "        elif self.scales != 1 and self.is_first_block == 1:\n",
    "            out = torch.cat((out, self.pool(x_scales[self.scales-1])), 1)\n",
    "\n",
    "        # 1*1卷积层\n",
    "        out = self.conv3(out)\n",
    "        out = self.bn3(out)\n",
    "\n",
    "        # 是否加入SE模块\n",
    "        if self.se is not None:\n",
    "            out = self.se(out)\n",
    "\n",
    "        # 添加triplet_attention\n",
    "\n",
    "        # 残差连接 out=F(X)+X\n",
    "        out += identity\n",
    "        #out = self.relu(out)\n",
    "        out = self.af(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "\n",
    "class Res2Net(nn.Module):\n",
    "    def __init__(self, block, layers, num_classes=1000, scales=4, groups=1, se=True):\n",
    "        super(Res2Net, self).__init__()\n",
    "        # 通道数初始化\n",
    "        self.inplanes = 64\n",
    "\n",
    "        # 起始：7*7的卷积层，3*3的最大池化层\n",
    "        self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)\n",
    "        self.bn1 = nn.BatchNorm2d(self.inplanes)\n",
    "        #self.relu = nn.ReLU(inplace=True)\n",
    "        if is_CELU :\n",
    "            self.af = nn.CELU(inplace=True)\n",
    "        elif is_ELU :\n",
    "            self.af = nn.ELU(inplace=True)\n",
    "        else :\n",
    "            self.af = nn.ReLU(inplace=True)\n",
    "        \n",
    "        \n",
    "        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "\n",
    "        # 残差结构\n",
    "        self.layer1 = self._make_layer(block, 64, layers[0], stride=1, scales=scales, groups=groups, se=se)\n",
    "        self.layer2 = self._make_layer(block, 128, layers[1], stride=2, scales=scales, groups=groups, se=se)\n",
    "        self.layer3 = self._make_layer(block, 256, layers[2], stride=2, scales=scales, groups=groups, se=se)\n",
    "        self.layer4 = self._make_layer(block, 512, layers[3], stride=2, scales=scales, groups=groups, se=se)\n",
    "\n",
    "        # 平均池化+全连接层\n",
    "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
    "        self.fc = nn.Linear(512 * block.expansion, num_classes)\n",
    "\n",
    "        # 权重初始化\n",
    "        for m in self.modules():\n",
    "            if isinstance(m, nn.Conv2d):\n",
    "                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "\n",
    "    def _make_layer(self, block, planes, layer, stride=1, scales=4, groups=1, se=True):\n",
    "        # 积步长不为1或深度扩张有变化，导致F(X)与X的shape不同的残差块，就要对X定义下采样函数，使之shape相同\n",
    "        downsample = None\n",
    "        if stride != 1 or self.inplanes != planes * block.expansion:\n",
    "            downsample = nn.Sequential(\n",
    "                nn.Conv2d(self.inplanes, planes * block.expansion, kernel_size=1, stride=stride, bias=False),\n",
    "                nn.BatchNorm2d(planes * block.expansion),\n",
    "            )\n",
    "\n",
    "        layers = []\n",
    "        # 第一个残差块需要下采样  def __init__(self, inplanes, planes, stride=1, downsample=None, scales=4, groups=1):\n",
    "        layers.append(block(self.inplanes, planes, stride=stride, downsample=downsample,\n",
    "                            scales=scales, groups=groups, is_first_block=1, se=se))\n",
    "        self.inplanes = planes * block.expansion\n",
    "\n",
    "        # 通过循环堆叠其余残差块\n",
    "        for i in range(1, layer):\n",
    "            layers.append(block(self.inplanes, planes, scales=scales, groups=groups, se=se))\n",
    "\n",
    "        return nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        #x = self.relu(x)\n",
    "        x = self.af(x)\n",
    "        x = self.maxpool(x)\n",
    "\n",
    "        x = self.layer1(x)\n",
    "        x = self.layer2(x)\n",
    "        x = self.layer3(x)\n",
    "        x = self.layer4(x)\n",
    "\n",
    "        x = self.avgpool(x)\n",
    "        x = torch.flatten(x, 1)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        return x\n",
    "\n",
    "\n",
    "# def __init__(self, block, layers, num_classes=1000, scales=4, groups=1):\n",
    "def res2net50(num_classes=100, scales=4, groups=1, se=True):\n",
    "    return Res2Net(Bottleneck, [3, 4, 6, 3], num_classes, scales, groups, se)\n",
    "\n",
    "def res2net18(num_classes=100, scales=4, groups=1, se=True):\n",
    "    return Res2Net(Bottleneck, [2, 2, 2, 2], num_classes, scales, groups, se)\n",
    "\n",
    "\n",
    "def res2net101(num_classes=100, scales=4, groups=1):\n",
    "    return Res2Net(Bottleneck, [3, 4, 23, 3], num_classes, scales, groups)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abcf60b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "from torch.optim import lr_scheduler\n",
    "#criterion = nn.CrossEntropyLoss().to(device)\n",
    "#optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)\n",
    "\n",
    "#optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001, momentum=0.9, weight_decay=5e-4)\n",
    "#optimizer = torch.optim.Adam(net.parameters(), lr=0.0004)\n",
    "\n",
    "# 定义损失函数和优化器\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "23014d4c",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def get_data_loader():    \n",
    "    transform_train = transforms.Compose([\n",
    "        transforms.RandomCrop(32, padding=4),\n",
    "        transforms.RandomHorizontalFlip(),\n",
    "        RandAugment(),  # 使用 RandAugment\n",
    "        transforms.ToTensor(),\n",
    "        #transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "        transforms.Normalize(\n",
    "            (0.5070751592371323, 0.48654887331495095, 0.4409178433670343),\n",
    "            (0.2673342858792401, 0.2564384629170883, 0.27615047132568404)\n",
    "        )\n",
    "    ])\n",
    "    \n",
    "    transform_test = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n",
    "    ])\n",
    "    \n",
    "\n",
    "    trainset = torchvision.datasets.CIFAR100(root='./data', train=True,\n",
    "                                             download=False, transform=transform_train)\n",
    "    trainloader = torch.utils.data.DataLoader(trainset, batch_size=128,\n",
    "                                              shuffle=True, num_workers=2)\n",
    "\n",
    "    testset = torchvision.datasets.CIFAR100(root='./data', train=F,\n",
    "                                            download=False, transform=transform_test)\n",
    "    testloader = torch.utils.data.DataLoader(testset, batch_size=128,\n",
    "                                             shuffle=False, num_workers=2)\n",
    "        # 实例化模型\n",
    "    return trainloader,testloader\n",
    "\n",
    "\n",
    "def plot_training_results(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs):\n",
    "    # Convert PyTorch tensors to NumPy arrays\n",
    "    train_losses = np.array(train_losses)\n",
    "    train_top1_accs = np.array(train_top1_accs)\n",
    "    train_top5_accs = np.array(train_top5_accs)\n",
    "    val_losses = np.array(val_losses)\n",
    "    val_top1_accs = np.array(val_top1_accs)\n",
    "    val_top5_accs = np.array(val_top5_accs)\n",
    "\n",
    "    plt.figure(figsize=(15, 5))\n",
    "    \n",
    "    # 绘制损失曲线\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(train_losses, label='Training Loss')\n",
    "    plt.plot(val_losses, label='Validation Loss')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Loss')\n",
    "    plt.title('Training/Validation Loss Curve')\n",
    "    plt.legend()\n",
    "\n",
    "    # 绘制准确率曲线\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(train_top1_accs, label='Training Top-1 Accuracy')\n",
    "    plt.plot(val_top1_accs, label='Validation Top-1 Accuracy')\n",
    "    plt.plot(train_top5_accs, label='Training Top-5 Accuracy')\n",
    "    plt.plot(val_top5_accs, label='Validation Top-5 Accuracy')\n",
    "    plt.xlabel('Epoch')\n",
    "    plt.ylabel('Accuracy')\n",
    "    plt.title('Training/Validation Accuracy Curve')\n",
    "    plt.legend()\n",
    "    \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "# 训练模型\n",
    "def train_model(model,trainloader,device,testloader, criterion, optimizer, scheduler, num_epochs=25):\n",
    "    since = time.time()\n",
    "       \n",
    "    train_losses = []\n",
    "    train_top1_accs = []\n",
    "    train_top5_accs = []\n",
    "    val_losses = []\n",
    "    val_top1_accs = []\n",
    "    val_top5_accs = []\n",
    "\n",
    "\n",
    "    best_model_wts = copy.deepcopy(model.state_dict())\n",
    "    best_acc = 0.0\n",
    " \n",
    "    for epoch in range(num_epochs):\n",
    "        print('第 {}/{} 轮训练'.format(epoch, num_epochs - 1))\n",
    "        print('-' * 10)\n",
    "\n",
    "        # 每个epoch都有训练和验证阶段\n",
    "        for phase in ['train', 'val']:\n",
    "            if phase == 'train':\n",
    "                model.train()  # 设置模型为训练模式\n",
    "            else:\n",
    "                model.eval()   # 设置模型为评估模式\n",
    "\n",
    "            running_loss = 0.0\n",
    "            running_corrects = 0\n",
    "            running_top1_corrects = 0\n",
    "            running_top5_corrects = 0\n",
    "\n",
    "            # 遍历数据\n",
    "            for inputs, labels in (trainloader if phase == 'train' else testloader):\n",
    "                inputs = inputs.to(device)\n",
    "                labels = labels.to(device)\n",
    "\n",
    "                # 零化参数梯度\n",
    "                optimizer.zero_grad()\n",
    "\n",
    "                # 前向传播\n",
    "                # 只在训练阶段追踪历史\n",
    "                with torch.set_grad_enabled(phase == 'train'):\n",
    "                    outputs = model(inputs)\n",
    "                    _, preds = torch.max(outputs, 1)\n",
    "                    loss = criterion(outputs, labels)\n",
    "\n",
    "                    # 计算 top-1 和 top-5 准确率\n",
    "                    _, top5_preds = torch.topk(outputs, 5, dim=1)\n",
    "                    top1_correct = torch.sum(preds == labels.data)\n",
    "                    top5_correct = torch.sum(top5_preds == labels.view(-1, 1))\n",
    "\n",
    "                    # 只有在训练阶段反向传播+优化\n",
    "                    if phase == 'train':\n",
    "                        loss.backward()\n",
    "                        optimizer.step()\n",
    "\n",
    "                # 统计\n",
    "                running_loss += loss.item() * inputs.size(0)\n",
    "                running_corrects += top1_correct\n",
    "                running_top1_corrects += top1_correct\n",
    "                running_top5_corrects += top5_correct\n",
    "\n",
    "            if phase == 'train':\n",
    "                scheduler.step()\n",
    "\n",
    "            epoch_loss = running_loss / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top1_acc = running_top1_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "            epoch_top5_acc = running_top5_corrects.double() / len(trainloader.dataset if phase == 'train' else testloader.dataset)\n",
    "\n",
    "            print('{} 损失: {:.4f} Top-1 准确率: {:.4f} Top-5 准确率: {:.4f}'.format(\n",
    "                phase, epoch_loss, epoch_top1_acc, epoch_top5_acc))\n",
    "            \n",
    "            # 保存损失和准确率\n",
    "            if phase == 'train':\n",
    "                train_losses.append(epoch_loss)\n",
    "                train_top1_accs.append(epoch_top1_acc)\n",
    "                train_top5_accs.append(epoch_top5_acc)\n",
    "            else:\n",
    "                val_losses.append(epoch_loss)\n",
    "                val_top1_accs.append(epoch_top1_acc)\n",
    "                val_top5_accs.append(epoch_top5_acc)\n",
    "\n",
    "            # 深度复制模型\n",
    "            if phase == 'val' and epoch_top1_acc > best_acc:\n",
    "                best_acc = epoch_top1_acc\n",
    "                best_model_wts = copy.deepcopy(model.state_dict())\n",
    "\n",
    "        print()\n",
    "\n",
    "    time_elapsed = time.time() - since\n",
    "    print('训练完成，耗时 {:.0f}m {:.0f}s'.format(\n",
    "        time_elapsed // 60, time_elapsed % 60))\n",
    "    print('最佳验证集准确率: {:4f}'.format(best_acc))\n",
    "\n",
    "    # 加载最佳模型权重\n",
    "    model.load_state_dict(best_model_wts)\n",
    "    \n",
    "    # 绘制训练结果曲线\n",
    "\n",
    "    return model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9dc2168a",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(device)\n",
    " ## 修改模型\n",
    "net = res2net18().to(device)\n",
    "num_epochs=50\n",
    " # 定义损失函数和优化器\n",
    "trainloader,testloader = get_data_loader()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(net.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)\n",
    "scheduler = lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.1)\n",
    "model,(train_losses, train_top1_accs, train_top5_accs, val_losses, val_top1_accs, val_top5_accs) = train_model(net,trainloader,device,testloader, criterion, optimizer, scheduler, num_epochs)  \n",
    "   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31661543",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "318a779b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8efdfd81",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eda4f864",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d4d9162",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
